#!/usr/bin/python
"""
Usage: python auto_anno.py annofile.txt seqfile.fasta
Author:			Jimmy Saw
Date modified:	07-03-2011

Description:    This program parses mate pairs with consistent taxonomic 
                assignments from phymmBL and extracts the sequences that match 
                the criteria, i.e., paired and both Gloeobacter.

Usage:          dissertation_PhymmBLParser.py <phymmbl result> <reads.fasta> > <out.fasta>

Example:        dissertation_PhymmBLParser.py results.03.xx.txt reads.fasta > taxa_pairs.fasta

Note:           Needs phymmBL output and sequence fasta file
"""
import sys
import re
from Bio import SeqIO

phymmblfile = sys.argv[1]
pf = open(phymmblfile, "rU")
pfl = pf.readlines()

seqfile = sys.argv[2]
seqs = SeqIO.parse(seqfile, "fasta")
seqs_dict = {}

for s in seqs:
    seqs_dict[s.id] = s

#print "Done loading sequences..."

left = re.compile('(\w+)_L')
right = re.compile('(\w+)_R')

gloeo = re.compile('Gloeobacter.*')

allreads = []
reads_dict = {}
readnames = []

for line in pfl[1:]:
    strand = ""
    l = line.split('\t')
    readname = l[0]
    bestmatch = l[1]
    score = float(l[2])
    genus = l[3]
    genusConf = float(l[4])
    family = l[5]
    familyConf = float(l[6])
    order = l[7]
    orderConf = float(l[8])
    clas = l[9]
    clsConf = float(l[10])
    phylum = l[11]
    phylumConf = float(l[12].rstrip())
    readnames.append(readname)
    reads_dict[readname] = ((bestmatch, score, genus, family, order, clas, phylum))

pairstokeep = []
singletons = []

for read in readnames:
    if left.match(read):
        pattern = left.match(read)
        mate_pair = pattern.group(1) + "_R"
        if read in reads_dict and mate_pair in reads_dict:
#            if reads_dict[read][2] == "Gloeobacter" and reads_dict[mate_pair][2] == "Gloeobacter": #both mating pairs must be Gloeobacter
            if reads_dict[read][2] == "Gloeobacter" or reads_dict[mate_pair][2] == "Gloeobacter": #at least one is Gloeobacter
                pairstokeep.append((read, mate_pair))
                readnames.remove(mate_pair) # remove the mate pair from list to avoid double-checking
#                print "Working on Left"
    elif right.match(read):
        pattern = right.match(read)
        mate_pair = pattern.group(1) + "_L"
        if read in reads_dict and mate_pair in reads_dict:
#            if reads_dict[read][2] == "Gloeobacter" and reads_dict[mate_pair][2] == "Gloeobacter": #both mating pairs must be Gloeobacter
            if reads_dict[read][2] == "Gloeobacter" or reads_dict[mate_pair][2] == "Gloeobacter": #at least one is Gloeobacter
                pairstokeep.append((read, mate_pair))
                readnames.remove(mate_pair) # remove the mate pair from list to avoid double-checking
#                print "Working on Right"
    else:
        if reads_dict[read][2] == "Gloeobacter":
            singletons.append(read)
#            print "Working on Singles"

print "Number of mate pairs with Gloeobacter assignment: ", len(pairstokeep)
print "Number of singletons with Gloeobacter assignment: ", len(singletons)
print "Total number of reads that are Gloeobacter: ", 2 * len(pairstokeep) + len(singletons)

for pair in pairstokeep:
    #if pair[0] in seqs_dict and pair[1] in seqs_dict:
    print ">" + pair[0]
    print seqs_dict[pair[0]].seq
    print ">" + pair[1]
    print seqs_dict[pair[1]].seq

for read in singletons:
    print ">" + read
    print seqs_dict[read].seq

pf.close()

